import sys
import geometry_msgs.msg
import moveit_commander
import moveit_msgs.msg
import rospy

class ArmManager():
    def __init__(self):
        moveit_commander.roscpp_initialize(sys.argv)

        # robot = moveit_commander.RobotCommander()
        # scene = moveit_commander.PlanningSceneInterface()
        group_name = "manipulator_i5"
        
        self.__move_group = moveit_commander.MoveGroupCommander(group_name)
        # planning_frame = self.__move_group.get_planning_frame()

    def get_current_pose(self, raw = False):
        if raw:
            return self.__move_group.get_current_pose()
        else:
            return self.__move_group.get_current_pose().pose  

    def get_current_joint_values(self):
        return self.__move_group.get_current_joint_values()

    def set_max_acceleration_and_velocity(self, acceleration=0.3, velocity=0.3):
        # 设置允许的最大速度和加速度
        self.__move_group.set_max_acceleration_scaling_factor(acceleration)
        self.__move_group.set_max_velocity_scaling_factor(velocity)
        self.__move_group.set_goal_position_tolerance(0.001)
        self.__move_group.set_goal_orientation_tolerance(0.01) 


    def back_to_zero(self):
        # 控制机械臂先回到初始化位置
        self.__move_group.set_named_target('zero')
        self.__move_group.go()
        rospy.sleep(1)
   
    def set_joint_values(self, j0=None, j1=None, j2=None, j3=None, j4=None, j5=None):
        joint_goal = self.__move_group.get_current_joint_values()
        joint_goal[0] = joint_goal[0] if j0 == None else j0
        joint_goal[1] = joint_goal[1] if j1 == None else j1
        joint_goal[2] = joint_goal[2] if j2 == None else j2
        joint_goal[3] = joint_goal[3] if j3 == None else j3
        joint_goal[4] = joint_goal[4] if j4 == None else j4
        joint_goal[5] = joint_goal[5] if j5 == None else j5
        
        self.__move_group.go(joint_goal, wait=True)
  
    def set_target_pose(self, position=None, orientation=None, flag=True):

        c_pose = self.__move_group.get_current_pose().pose  

        # self.back_to_zero()  
         
        # 当运动规划失败后，允许重新规划
        # self.__move_group.allow_replanning(True)
        
        # 设置机械臂工作空间中的目标位姿，位置使用x、y、z坐标描述，
        # 姿态使用四元数描述，基于base_link坐标系
        target_pose = geometry_msgs.msg.Pose()

        #末端姿态，四元数
        target_pose.orientation.x = c_pose.orientation.x if orientation is None else orientation[0]
        target_pose.orientation.y = c_pose.orientation.y if orientation is None else orientation[1]
        target_pose.orientation.z = c_pose.orientation.z if orientation is None else orientation[2]
        target_pose.orientation.w = c_pose.orientation.w if orientation is None else orientation[3]
        
        #末端位置  
        target_pose.position.x = c_pose.position.x if position is None else position[0]
        target_pose.position.y = c_pose.position.y if position is None else position[1]
        target_pose.position.z = c_pose.position.z if position is None else position[2]

        # 设置机器臂当前的状态作为运动初始状态
        self.__move_group.set_start_state_to_current_state()
        
        # 设置机械臂终端运动的目标位姿
        self.__move_group.set_pose_target(target_pose)
        
        # 规划运动路径，返回虚影的效果
        traj = self.__move_group.plan()
        
        # 按照规划的运动路径控制机械臂运动
        self.__move_group.execute(traj,flag)
        rospy.sleep(1)  #执行完成后休息1s)
 
        # # 关闭并退出moveit
        # moveit_commander.roscpp_shutdown()
        # moveit_commander.os._exit(0)

    def stop(self):
        self.__move_group.stop()
        self.__move_group.clear_pose_targets()


    

if __name__ == "__main__":
    rospy.init_node('ARM_MANAGER', anonymous=False)

    am = ArmManager()

    print(am.get_current_pose())

    am.set_target_pose(j5=1.57)
    
